#
# Copyright (c) 2011-2012 Johann Hanne
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
#


import OMFParser

class OMFObject(object):

    class segment(object):
        def __init__(self):
            self.group = None

    class segmentgroup(object):
        def __init__(self):
            pass

    def __init__(self,fname):
        f = open(fname, 'rb')

        omf = f.read()
        omflen = len(omf)

        self.filename = fname
        self.recordlist = []

        offset = 0
        while offset < omflen:
            obj = OMFParser.get_next_record(omf, offset)
            self.recordlist.append(obj)
            offset += obj.reclen + 3

        f.close()

        self.fixupp_thread_target = [None, None, None, None]
        self.fixupp_thread_frame  = [None, None, None, None]

        self.fixablechunk = None

    def process(self, out):
        # ordered collections
        self.names = []
        self.segments = []
        self.segmentgroups = []
        self.extnames = []

        for i, record in enumerate(self.recordlist):
            if isinstance(record, OMFParser.OMF_LNAMES):
                for entry in record.entries:
                    self.names.append(entry.strcontent)
            elif isinstance(record, OMFParser.OMF_SEGDEF):
                if record.segment_name_index < 1 or record.segment_name_index > len(self.names):
                    raise RuntimeError("SEGDEF record (#{0}) contains invalid segment_name_index {1}".format(i, record.segment_name_index))
                if record.class_name_index < 1 or record.class_name_index > len(self.names):
                    raise RuntimeError("SEGDEF record (#{0}) contains invalid class_name_index {1}".format(i, record.class_name_index))
                if record.overlay_name_index < 1 or record.overlay_name_index > len(self.names):
                    raise RuntimeError("SEGDEF record (#{0}) contains invalid overlay_name_index {1}".format(i, record.overlay_name_index))

                segment = self.segment()

                segment.segattr_alignment = record.segattr_alignment
                segment.segattr_combination = record.segattr_combination
                segment.segattr_big = record.segattr_big
                segment.segattr_p = record.segattr_p
                if segment.segattr_alignment == 0:
                    segment.segattr_frame_number = record.segattr_frame_number
                    segment.segattr_offset = record.segattr_offset
                segment.segment_length = record.segment_length
                segment.segment_name = self.names[record.segment_name_index - 1]
                segment.class_name = self.names[record.class_name_index - 1]
                segment.overlay_name = self.names[record.overlay_name_index - 1]

                if segment.overlay_name != '':
                    raise RuntimeError("SEGDEF record (#{0}) contains non-empty overlay name '{1}' which is unsupported".format(i, record.overlay_name))

                segment.chunk = out.newsegment(self.filename, segment.segment_name, segment.class_name, segment.segment_length, segment.segattr_alignment, segment.segattr_combination)

                self.segments.append(segment)
            elif isinstance(record, OMFParser.OMF_GRPDEF):
                if record.group_name_index < 1 or record.group_name_index > len(self.names):
                    raise RuntimeError("GRPDEF record (#{0}) contains invalid group_name_index {1}".format(i, record.group_name_index))

                if len(record.entries) == 0:
                    raise RuntimeError("GRPDEF record (#{0}) contains no entries?!".format(i))

                groupmembers = []

                for entry in record.entries:
                    if entry.type != 0xff:
                        raise RuntimeError("GRPDEF record (#{0}) contains entry with invalid type {1}".format(i, entry.type))
                    if entry.segment_definition < 1 or entry.segment_definition > len(self.segments):
                        raise RuntimeError("GRPDEF record (#{0}) contains entry with invalid segment_definition {1}".format(i, entry.segment_definition))
                    if self.segments[entry.segment_definition - 1].group is not None:
                        raise RuntimeError("GRPDEF record (#{0}) contains group for segment_definition {1} which already is assigned to a group".format(i, entry.segment_definition))
                    self.segments[entry.segment_definition - 1].group = self.names[record.group_name_index - 1]

                    groupmembers.append(self.segments[entry.segment_definition - 1].chunk)

                segmentgroup = self.segmentgroup()
                segmentgroup.chunkgroup = out.newsegmentgroup(self.filename, self.names[record.group_name_index - 1], groupmembers)
                self.segmentgroups.append(segmentgroup)

            elif isinstance(record, (OMFParser.OMF_PUBDEF, OMFParser.OMF_LPUBDEF)):
                if record.base_segment_index < 1 or record.base_segment_index > len(self.segments):
                    raise RuntimeError("(L)PUBDEF record (#{0}) contains invalid base_segment_index {1}".format(i, record.base_segment_index))
                if record.base_group_index > 0 and record.base_group_index > len(self.segmentgroups):
                    raise RuntimeError("(L)PUBDEF record (#{0}) contains invalid base_group_index {1}".format(i, record.base_group_index))
                if record.base_frame != 0:
                    raise RuntimeError("(L)PUBDEF record (#{0}) contains base_frame != 0 which is unsupported".format(i))

                # TODO: There seems to be a bit of redundant information in
                # PUBDEF:it has both base_group_index and base_segment_index,
                # but GRPDEF records already define the group membership of
                # segments; may be we should check if the information is
                # consistent here...

                for entry in record.entries:
                    if isinstance(record, OMFParser.OMF_PUBDEF):
                        publicname = entry.strcontent
                    else:
                        # i.e. LPUBDEF
                        publicname = self.filename + "$" + entry.strcontent

                    self.segments[record.base_segment_index - 1].chunk.register_public_name(publicname, entry.public_offset)

            elif isinstance(record, OMFParser.OMF_LEDATA):
                if record.segment_index < 1 or record.segment_index > len(self.segments):
                    raise RuntimeError("LEDATA record (#{0}) contains invalid segment_index {1}".format(i, record.segment_index))

                self.fixablechunk = self.segments[record.segment_index - 1].chunk.setdata(record.data_offset, record.data, False)

            elif isinstance(record, OMFParser.OMF_LIDATA):
                if record.segment_index < 1 or record.segment_index > len(self.segments):
                    raise RuntimeError("LIDATA record (#{0}) contains invalid segment_index {1}".format(i, record.segment_index))

                self.fixablechunk = self.segments[record.segment_index - 1].chunk.setdata(record.data_offset, record.data, True)

            elif isinstance(record, OMFParser.OMF_EXTDEF):
                for entry in record.entries:
                    self.extnames.append(entry.strcontent)

            elif isinstance(record, OMFParser.OMF_LEXTDEF):
                for entry in record.entries:
                    self.extnames.append(self.filename + "$" + entry.strcontent)

            elif isinstance(record, OMFParser.OMF_FIXUPP):
                for subrecord in record.subrecords:
                    if isinstance(subrecord, OMFParser.OMF_FIXUPP.THREAD_subrecord):

                        if subrecord.thread_method >= 0 and subrecord.thread_method <= 2:
                            index = subrecord.index
                        else:
                            index = None

                        if subrecord.thread_d == 0:
                            self.fixupp_thread_target[subrecord.thread_thred] = (subrecord.thread_method, index)
                        else:
                            self.fixupp_thread_frame[subrecord.thread_thred]  = (subrecord.thread_method, index)

                    elif isinstance(subrecord, OMFParser.OMF_FIXUPP.FIXUP_subrecord):
                        if self.fixablechunk is None:
                            raise RuntimeError("FIXUPP record (#{0}) contains FIXUP subrecord but is not preceded by fixable data".format(i))

                        if subrecord.fix_data_f == 1:
                            if self.fixupp_thread_frame[subrecord.fix_data_frame] is None:
                                raise RuntimeError("FIXUP subrecord refers to undefined frame thread {0}".format(subrecord.fix_data_frame))
                            (frame_method, frame_datum) = self.fixupp_thread_frame[subrecord.fix_data_frame]
                        else:
                            frame_method = subrecord.fix_data_frame
                            if frame_method >= 0 and frame_method <= 2:
                                frame_datum = subrecord.frame_datum
                            else:
                                frame_datum = None

                        if subrecord.fix_data_t == 1:
                            if self.fixupp_thread_target[subrecord.fix_data_target] is None:
                                raise RuntimeError("FIXUP subrecord refers to undefined target thread {0}".format(subrecord.fix_data_target))
                            (target_method, target_datum) = self.fixupp_thread_target[subrecord.fix_data_targt]
                        else:
                            target_method = subrecord.fix_data_targt
                            if target_method >= 0 and target_method <= 2:
                                target_datum = subrecord.target_datum
                            else:
                                target_datum = None

                        target_method |= (subrecord.fix_data_p << 2)

                        #if target_method == 0:
                        #    # Specified by a SEGDEF index.
                        #elif target_method == 1:
                        #    # Specified by a GRPDEF index.
                        #elif target_method == 2:
                        #    # Specified by a EXTDEF index.
                        #elif target_method == 3:
                        #    # Specified by an explicit frame number (not supported by Microsoft LINK or IBM LINK386).
                        if target_method == 4:
                            # Specified by a SEGDEF index only; the displacement in the FIXUP subrecord is assumed to be 0.
                            if target_datum < 1 or target_datum > len(self.segments):
                                raise RuntimeError("FIXUP subrecord contains invalid segment index {0}".format(target_datum))
                            target_datum_resolved = self.segments[target_datum - 1].chunk
                            target_displacement = 0
                        #elif target_method == 5:
                        #    # Specified by a GRPDEF index only; the displacement in the FIXUP subrecord is assumed to be 0.
                        elif target_method == 6:
                            # Specified by a EXTDEF index only; the displacement in the FIXUP subrecord is assumed to be 0.
                            if target_datum < 1 or target_datum > len(self.extnames):
                                raise RuntimeError("FIXUP subrecord contains invalid external name index {0}".format(target_datum))
                            target_datum_resolved = self.extnames[target_datum - 1]
                            target_displacement = 0
                        else:
                            raise RuntimeError("FIXUP subrecord uses target method {0} which is not (yet?) supported...".format(target_method))

                        if frame_method == 0:
                            # The FRAME is specified by a SEGDEF index.
                            if frame_datum < 1 or frame_datum > len(self.segments):
                                raise RuntimeError("FIXUP subrecord contains invalid segment index {0}".format(frame_datum))
                            frame_datum_resolved = self.segments[frame_datum - 1].chunk
                        elif frame_method == 1:
                            # The FRAME is specified by a GRPDEF index.
                            if frame_datum < 1 or frame_datum > len(self.segmentgroups):
                                raise RuntimeError("FIXUP subrecord contains invalid segment group index {0}".format(frame_datum))
                            frame_datum_resolved = self.segmentgroups[frame_datum - 1].chunkgroup
                        elif frame_method == 5:
                            # The FRAME is determined by the TARGET's segment, group, or external index.
                            frame_datum_resolved = target_datum_resolved
                        else:
                            raise RuntimeError("FIXUP subrecord uses frame method {0} which is not (yet?) supported...".format(frame_method))

                        self.fixablechunk.addfixup(m=subrecord.fixup_m, location=subrecord.fixup_location,
                            data_record_offset=subrecord.fixup_data_record_offset,
                            frame_method=frame_method, frame_datum=frame_datum_resolved,
                            target_method=target_method, target_datum=target_datum_resolved,
                            target_displacement=target_displacement)

                    else:
                        # This would be an internal error...
                        raise RuntimeError("FIXUPP record (#{0}) contains invalid subrecord {1}".format(i, type(subrecord)))

        #for segment in self.segments:
        #    print("Segment: name '{0}', class '{1}', group '{2}', length: {3}".format(segment.segment_name, segment.class_name, segment.group, segment.segment_length))

